// Copyright 2018 The Cockroach Authors.
// Copyright (c) 2022-present, Shanghai Yunxi Technology Co, Ltd. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// This software (KWDB) is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"os"
	"path/filepath"
	"regexp"

	"gitee.com/kwbasedb/kwbase/pkg/sql/colexec/execerror"
	"github.com/pkg/errors"
)

var (
	errInvalidArgCount = errors.New("invalid number of arguments")
)

func main() {
	gen := execgen{stdErr: os.Stderr}
	if !gen.run(os.Args[1:]...) {
		os.Exit(2)
	}
}

type execgen struct {
	// useGoFmt runs the go fmt tool on code generated by execgen, if this setting
	// is true.
	useGoFmt bool

	// stdErr is the writer to which all standard error output will be redirected.
	stdErr io.Writer

	// cmdLine stores the set of flags used to invoke the Execgen tool.
	cmdLine *flag.FlagSet
}

type generator func(io.Writer) error

var generators = make(map[string]entry)

type entry struct {
	fn  generator
	dep string
}

func registerGenerator(g generator, filename, dep string) {
	if _, ok := generators[filename]; ok {
		execerror.VectorizedInternalPanic(fmt.Sprintf("%s generator already registered", filename))
	}
	generators[filename] = entry{fn: g, dep: dep}
}

func (g *execgen) run(args ...string) bool {
	// Parse command line.
	var printDeps bool
	g.cmdLine = flag.NewFlagSet("execgen", flag.ContinueOnError)
	g.cmdLine.SetOutput(g.stdErr)
	g.cmdLine.Usage = g.usage
	g.cmdLine.BoolVar(&g.useGoFmt, "useGoFmt", true, "run go fmt on generated code")
	g.cmdLine.BoolVar(&printDeps, "M", false, "print the dependency list")
	err := g.cmdLine.Parse(args)
	if err != nil {
		return false
	}

	// Get remaining args after any flags have been parsed.
	args = g.cmdLine.Args()
	if len(args) < 1 {
		g.cmdLine.Usage()
		g.reportError(errInvalidArgCount)
		return false
	}

	for _, out := range args {
		_, file := filepath.Split(out)
		entry := generators[file]
		if entry.fn == nil {
			g.reportError(errors.Errorf("unrecognized filename: %s", file))
			return false
		}
		if printDeps {
			if entry.dep == "" {
				// This output file has no template dependency (its template
				// is embedded entirely in execgen). Skip it.
				continue
			}
			fmt.Printf("%s: %s\n", out, entry.dep)
		} else {
			if err := g.generate(entry.fn); err != nil {
				g.reportError(err)
				return false
			}
		}
	}

	return true
}

var emptyCommentRegex = regexp.MustCompile(`[ \t]*//[ \t]*\n`)
var emptyBlockCommentRegex = regexp.MustCompile(`[ \t]*/\*[ \t]*\*/[ \t]*\n`)

func (g *execgen) generate(genFunc generator) error {
	var buf bytes.Buffer
	buf.WriteString("// Code generated by execgen; DO NOT EDIT.\n")

	err := genFunc(&buf)
	if err != nil {
		return err
	}

	b := buf.Bytes()
	// Delete empty comments ( // or /* */) that tend to get generated by templating.
	b = emptyCommentRegex.ReplaceAll(b, []byte{})
	b = emptyBlockCommentRegex.ReplaceAll(b, []byte{})

	if g.useGoFmt {
		oldB := b
		b, err = format.Source(b)
		if err != nil {
			// Write out incorrect source for easier debugging.
			b = oldB
			err = errors.Wrap(err, "Code formatting failed with Go parse error")
		}
	} else {
		b = buf.Bytes()
	}

	// Ignore any write error if another error already occurred.
	_, writeErr := os.Stdout.Write(b)
	if err != nil {
		return err
	}
	return writeErr
}

// usage is a replacement usage function for the flags package.
func (g *execgen) usage() {
	fmt.Fprintf(g.stdErr, "Execgen is a tool for generating templated code related to ")
	fmt.Fprintf(g.stdErr, "columnarized execution.\n\n")

	fmt.Fprintf(g.stdErr, "Usage:\n")
	fmt.Fprintf(g.stdErr, "\texecgen [path]...\n\n")

	fmt.Fprintf(g.stdErr, "Supported filenames are:\n")
	for filename := range generators {
		fmt.Fprintf(g.stdErr, "\t%s\n", filename)
	}
	fmt.Fprintf(g.stdErr, "\n")

	fmt.Fprintf(g.stdErr, "Flags:\n")
	g.cmdLine.PrintDefaults()
	fmt.Fprintf(g.stdErr, "\n")
}

func (g *execgen) reportError(err error) {
	fmt.Fprintf(g.stdErr, "ERROR: %v\n", err)
}
